package MAIN;
import java.io.BufferedInputStream;
import java.io.BufferedOutputStream;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.InputStream;
import java.io.ObjectInputStream;
import java.io.ObjectOutputStream;
import java.math.BigInteger;
import java.nio.charset.Charset;
import java.security.KeyFactory;
import java.security.KeyPair;
import java.security.KeyPairGenerator;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.security.spec.RSAPrivateKeySpec;
import java.security.spec.RSAPublicKeySpec;
import javax.crypto.Cipher;

//http://www.javamex.com/tutorials/cryptography/rsa_encryption.shtml

public class usc_RSAKeyPair 
{
    /*
     * the public key consists of the "modulus" and a "public exponent";
     * the private key consists of that same "modulus" plus a "private exponent";
     */
    public static String PUBLIC_KEY_FILE = "";
    public static String PRIVATE_KEY_FILE = "";
    
    public static void generate() throws Exception
    {
         KeyPairGenerator kpg = KeyPairGenerator.getInstance("RSA");
         kpg.initialize(2048);
         KeyPair kp = kpg.genKeyPair();
         KeyFactory fact = KeyFactory.getInstance("RSA");
         RSAPublicKeySpec pub = fact.getKeySpec(kp.getPublic(),RSAPublicKeySpec.class);
         saveToFile(PUBLIC_KEY_FILE, pub.getModulus(), pub.getPublicExponent());
         //System.out.println("public exponent: " + pub.getPublicExponent());
         RSAPrivateKeySpec priv = fact.getKeySpec(kp.getPrivate(), RSAPrivateKeySpec.class);
         saveToFile(PRIVATE_KEY_FILE, priv.getModulus(), priv.getPrivateExponent());
         //System.out.println("private exponent: " + priv.getPrivateExponent());
    }
    
    private static void saveToFile(String fileName,
                               BigInteger mod, BigInteger exp) throws Exception 
    {
        ObjectOutputStream oout = new ObjectOutputStream(
                new BufferedOutputStream(new FileOutputStream(fileName)));
        try {
            oout.writeObject(mod);
            oout.writeObject(exp);
        } catch (Exception e) {
            throw new Exception(e);
        } finally {
            oout.flush();
            oout.close();
        }
         //System.out.println("exponent: " + exp.toString());
    }
    
    public static PublicKey readPublicKey() throws Exception 
    {
        InputStream in = new FileInputStream(PUBLIC_KEY_FILE);
        ObjectInputStream oin =
                new ObjectInputStream(new BufferedInputStream(in));
        try {
            BigInteger m = (BigInteger) oin.readObject();
            BigInteger e = (BigInteger) oin.readObject();
            RSAPublicKeySpec keySpec = new RSAPublicKeySpec(m, e);
            KeyFactory fact = KeyFactory.getInstance("RSA");
            PublicKey pubKey = fact.generatePublic(keySpec);
            return pubKey;
        } catch (Exception e) {
            throw new Exception(e);
        } finally {
            oin.close();
        }
    }
    
    public static PrivateKey readPrivateKey() throws Exception 
    {
        InputStream in = new FileInputStream(PRIVATE_KEY_FILE);
        ObjectInputStream oin =
                new ObjectInputStream(new BufferedInputStream(in));
        try {
            BigInteger m = (BigInteger) oin.readObject();
            BigInteger e = (BigInteger) oin.readObject();
            RSAPrivateKeySpec keySpec = new RSAPrivateKeySpec(m, e);
            KeyFactory fact = KeyFactory.getInstance("RSA");
            PrivateKey privKey = fact.generatePrivate(keySpec);
            return privKey;
        } catch (Exception e) {
            throw new Exception(e);
        } finally {
            oin.close();
        }
    }
    
    
    public static byte[] PAIR1_Encrypt(byte[] input, PrivateKey privKey) throws Exception
    {
        byte[] output = null;
        Cipher cipher = Cipher.getInstance("RSA");
        cipher.init(Cipher.ENCRYPT_MODE, privKey);
        output = cipher.doFinal(input);
        return output;
    }
    public static byte[] PAIR1_Decrypt(byte[] input, PublicKey pubKey) throws Exception
    {
        byte[] output = null;
        Cipher cipher = Cipher.getInstance("RSA");
        cipher.init(Cipher.DECRYPT_MODE, pubKey);
        output = cipher.doFinal(input);
        return output;
    }

        public static byte[] PAIR2_Encrypt(byte[] input, PublicKey pubKey) throws Exception
    {
        byte[] output = null;
        Cipher cipher = Cipher.getInstance("RSA");
        cipher.init(Cipher.ENCRYPT_MODE, pubKey);
        output = cipher.doFinal(input);
        return output;
    }
    public static byte[] PAIR2_Decrypt(byte[] input, PrivateKey privKey) throws Exception
    {
        byte[] output = null;
        Cipher cipher = Cipher.getInstance("RSA");
        cipher.init(Cipher.DECRYPT_MODE, privKey);
        output = cipher.doFinal(input);
        return output;
    }
    
    //=====================================================
    public static void main(String[] args) throws Exception 
    {
        
        usc_RSAKeyPair.PUBLIC_KEY_FILE = "resource/publickey_server";
        usc_RSAKeyPair.PRIVATE_KEY_FILE = "resource/privatekey_server";
        //usc_RSAKeyPair.generate();
        PublicKey pubKey = usc_RSAKeyPair.readPublicKey();    
        PrivateKey privKey = usc_RSAKeyPair.readPrivateKey();    
        System.out.println(pubKey);
        System.out.println("================");
        System.out.println(privKey);
        
        //~~~~~~Encrypt and Decrypt~~~~~~~
        System.out.println("--- --- --- --- --- --- --- --- ---");
        String s1 = "Anh yêu em";
        System.out.println("User input: " + s1);
        byte[] s1_b = s1.getBytes(Charset.forName("UTF-8"));
        //byte[] s1_b_encrypted = PAIR1_Encrypt(s1_b, privKey);
        byte[] s1_b_encrypted = PAIR2_Encrypt(s1_b, pubKey);
        String s1_b_encrypted_out = new String(s1_b_encrypted,"UTF-8");
        System.out.println("Encrypted output: " + s1_b_encrypted_out);
        
        //byte[] s1_b_decrypted = PAIR1_Decrypt(s1_b_encrypted, pubKey);
        byte[] s1_b_decrypted = PAIR2_Decrypt(s1_b_encrypted, privKey);
        String s1_b_decrypted_out = new String(s1_b_decrypted,"UTF-8");
        System.out.println("Decrypted output: " + s1_b_decrypted_out);
    }
    
}
